#!/bin/env python
# -*- coding:utf-8
"""
@author:xingz
@file:Connectfour.py
@time:2021/11/17
"""
from __future__ import annotations
from typing import List, Optional, Tuple
from enum import Enum
from board import Board, Piece, Move
import copy


class C4Piece(Piece, Enum):
    B = "B"
    R = "R"
    E = " "  # stand-in for empty

    @property
    def opposite(self) -> Piece:
        if self == C4Piece.B:
            return C4Piece.R
        elif self == C4Piece.R:
            return C4Piece.B
        else:
            return C4Piece.E

    def __str__(self):
        return self.value


def generate_segments(num_columns: int, num_rows: int, segment_length: int) -> List[List[Tuple[int, int]]]:
    segments: List[List[Tuple[int, int]]] = []
    for c in range(num_columns):
        for r in range(num_rows - segment_length + 1):
            segment = []
            for t in range(segment_length):
                segment.append((c, r + t))
            segments.append(segment)

    for c in range(num_columns - segment_length + 1):
        for r in range(num_rows):
            segment = []
            for t in range(segment_length):
                segment.append((c + t, r))
            segments.append(segment)

    for c in range(num_columns - segment_length + 1):
        for r in range(num_rows - segment_length + 1):
            segment = []
            for t in range(segment_length):
                segment.append((c + t, r + t))
            segments.append(segment)

    for c in range(0, num_columns - segment_length + 1, -1):
        for r in range(segment_length - 1, num_rows):
            segment = []
            for t in range(segment_length):
                segment.append((c + t, r - t))
            segments.append(segment)
    return segments


class C4Board(Board):
    NUM_ROWS: int = 6
    NUM_COLUMNS: int = 7
    SEGMENT_LENGTH: int = 4
    SEGMENTS: List[List[Tuple[int, int]]] = generate_segments(NUM_COLUMNS, NUM_ROWS, SEGMENT_LENGTH)

    def __init__(self, position: Optional[List[C4Board.Column]] = None, turn: C4Piece = C4Piece.B):
        if position is None:
            self.position = [C4Board.Column() for i in range(self.NUM_COLUMNS)]
        else:
            self.position = position
        self._turn: C4Piece = turn

    @property
    def turn(self) -> Piece:
        return self._turn

    def move(self, move: Move) -> Board:
        temp_position = self.position.copy()
        for c in range(C4Board.NUM_COLUMNS):
            temp_position[c] = self.position[c].copy()
        temp_position[move].push(self._turn)
        return C4Board(temp_position, self._turn.opposite)

    @property
    def legal_moves(self) -> List[Move]:
        return [Move(c) for c in range(C4Board.NUM_COLUMNS) if not self.position[c].full]

    def _count_segment(self, segment: List[Tuple[int, int]]) -> Tuple[int, int]:
        black_count: int = 0
        red_count: int = 0
        for column, row in segment:
            if self.position[column][row] == C4Piece.B:
                black_count += 1
            elif self.position[column][row] == C4Piece.R:
                red_count += 1
        return black_count, red_count

    @property
    def is_win(self) -> bool:
        for segment in C4Board.SEGMENTS:
            black_count, red_count = self._count_segment(segment)
            if black_count == C4Board.SEGMENT_LENGTH or red_count == C4Board.SEGMENT_LENGTH:
                return True
        return False

    def _evaluate_segment(self, segment: List[Tuple[int, int]], player: Piece) -> float:
        black_count, red_count = self._count_segment(segment)
        if red_count > 0 and black_count > 0:
            return 0
        count: int = max(red_count, black_count)
        score: float = 0
        if count == 2:
            score = 1
        elif count == 3:
            score = 100
        elif count == 4:
            score = 1000000
        color: C4Piece = C4Piece.B
        if red_count > red_count:
            color = C4Piece.R
        if color != player:
            return -score
        return score

    def evaluate(self, player: Piece) -> float:
        total = 0
        for segments in C4Board.SEGMENTS:
            total += self._evaluate_segment(segments, player)
        return total

    def __repr__(self):
        displayer: str = ""
        for r in reversed(range(C4Board.NUM_ROWS)):
            displayer += "|"
            for c in range(C4Board.NUM_COLUMNS):
                displayer += f"{self.position[c][r]}" + "|"
            displayer += "\n"
        return displayer

    class Column:
        def __init__(self):
            self._container: List[C4Piece] = []

        @property
        def full(self) -> bool:
            return len(self._container) == C4Board.NUM_ROWS

        def push(self, item: C4Piece) -> None:
            if self.full:
                raise OverflowError("Try to push piece to ful column")
            self._container.append(item)

        def __getitem__(self, item) -> C4Piece:
            if item > len(self._container) - 1:
                return C4Piece.E
            return self._container[item]

        def __repr__(self):
            return repr(self._container)

        def copy(self) -> C4Board.Column:
            temp: C4Board.Column = C4Board.Column()
            temp._container = self._container.copy()
            return temp


class C4Board2(Board):
    NUM_COLUMNS = 7
    NUM_ROWS = 7
    SEGMENT_LENGTH = 4
    SEGMENTS: List[List[Tuple, Tuple]] = generate_segments(NUM_COLUMNS, NUM_ROWS, SEGMENT_LENGTH)

    def __init__(self, position: List[List[Tuple[int, int]]] = None, turn: C4Piece = C4Piece.B):
        if position is None:
            position = self._init_position()
        self.position = position
        self._turn = turn

    def _init_position(self):
        position = []
        row = [C4Piece.E] * self.NUM_ROWS
        for c in range(self.NUM_COLUMNS):
            position.append(copy.deepcopy(row))
        return position

    @property
    def turn(self) -> Piece:
        return self._turn

    def move(self, move: Tuple[Move, Move]) -> Board:
        from copy import deepcopy
        tmp_position = deepcopy(self.position)
        tmp_position[move[0]][move[1]] = self._turn
        return C4Board2(tmp_position, self._turn.opposite)

    @property
    def legal_moves(self) -> List[Tuple[Move, Move]]:
        return [(Move(c), Move(r)) for c in range(len(self.position)) for r in range(len(self.position[0])) if
                self.position[c][r] == C4Piece.E]

    def _count_segment(self, segment: List[Tuple[int, int]]) -> Tuple[int, int]:
        black_count: int = 0
        red_count: int = 0
        for column, row in segment:
            if self.position[column][row] == C4Piece.B:
                black_count += 1
            elif self.position[column][row] == C4Piece.R:
                red_count += 1
        return black_count, red_count

    @property
    def is_win(self) -> bool:
        for segment in self.SEGMENTS:
            black_count, red_count = self._count_segment(segment)
            if black_count == C4Board2.SEGMENT_LENGTH or red_count == C4Board2.SEGMENT_LENGTH:
                return True
        return False

    def _evaluate_segment(self, segment: List[Tuple[int, int]], player: Piece) -> float:
        black_count, red_count = self._count_segment(segment)
        if red_count > 0 and black_count > 0:
            return 0
        count: int = max(red_count, black_count)
        score: float = 0
        if count == 2:
            score = 1
        elif count == 3:
            score = 100
        elif count == 4:
            score = 1000000
        color: C4Piece = C4Piece.B
        if red_count > red_count:
            color = C4Piece.R
        if color != player:
            return -score
        return score

    def evaluate(self, player: Piece) -> float:
        total = 0
        for segments in C4Board.SEGMENTS:
            total += self._evaluate_segment(segments, player)
        return total

    def __repr__(self):
        displayer: str = ""
        for r in reversed(range(C4Board2.NUM_ROWS)):
            displayer += "|"
            for c in range(C4Board2.NUM_COLUMNS):
                displayer += f"{self.position[c][r]}" + "|"
            displayer += "\n"
        return displayer


if __name__ == '__main__':
    b = C4Board2()
    b = b.move((Move(1), Move(2)))
    print(b)
